#include "OptimiseRt.h"


double computeEnergyLM(vector<MarkerPixel> &M, const superFrame &I, IntrinsicMatrix&K, Camera&C, const float * p, const vector<Point3f> &kp3, const vector<Point2f> &kp2, double w1, double w2)
{
	vector<Point3f> kp3_c;
	C.TransformToCameraPlane(kp3, kp3_c);
	vector<Point2f> kp2_d;
	K.projectToImage(kp3_c, kp2_d);
	double e = 0, e2 = 0;
	for (size_t i = 0; i < kp2_d.size(); i++)
	{
		e += normL2Sqr(kp2_d[i], kp2[i]);
	}
	e *= w2;

	const int m_x = I.cols;
	const int m_y = I.rows;
	Point3f tt;
	Point2i ttt;
	Mat hah = Mat_<uchar>(m_y, m_x, (uchar)0);

	for (size_t i = 0; i < M.size(); i++)
	{
		auto &m = M[i];
		C.TransformToCameraPlane(m.m_3d, tt);
		K.projectToImage(tt, ttt);
		if (ttt.x < 0 || ttt.x >= m_x || ttt.y < 0 || ttt.y >= m_y)
			continue;
		hah.ptr<uchar>(ttt.y)[ttt.x] = m.i;
		e2 += pow(p[0] * I.ptr<FramePixel>(ttt.y)[ttt.x].i + p[1] - m.i, 2);
	}
	e2 *= w1;
	return e2 + e;
}


void computeGradient(const Mat &_src, Mat &output)
{
	Mat src;
	if (_src.type() == CV_8UC3)
	{
		cvtColor(_src, src, CV_BGR2GRAY);
	}
	else
	{
		src = _src;
	}
	output = Mat_<Vec3f>(_src.size());
	for (size_t i = 1; i < src.rows - 1; i++)
	{
		auto p0 = src.ptr<uchar>(i - 1);
		auto p1 = src.ptr<uchar>(i);
		auto p2 = src.ptr<uchar>(i + 1);
		auto g = output.ptr<FramePixel>(i);
		for (size_t j = 0; j < src.cols; j++)
		{
			g[j].i = p1[j];
			g[j].gy = (p2[j] - p0[j]) / 2;
		}
	}
	auto g = output.ptr<FramePixel>(0);
	auto g2 = output.ptr<FramePixel>(output.rows - 1);
	auto p0 = src.ptr<uchar>(0);
	auto p1 = src.ptr<uchar>(1);
	auto p2 = src.ptr<uchar>(src.rows - 2);
	auto p3 = src.ptr<uchar>(src.rows - 1);
	for (size_t j = 0; j < src.cols; j++)
	{
		g[j].gy = (p1[j] - p0[j]) / 2;
		g2[j].gy = (p3[j] - p2[j]) / 2;
		g[j].i = p0[j];
		g2[j].i = p3[j];
	}

	for (size_t i = 0; i < src.rows; i++)
	{
		auto p0 = src.ptr<uchar>(i);
		auto g = output.ptr<FramePixel>(i);
		for (size_t j = 1; j < src.cols - 1; j++)
		{
			g[j].gx = (p0[j + 1] - p0[j - 1]) / 2;
		}
	}

	for (size_t i = 0; i < src.rows; i++)
	{
		const int cols = output.cols;
		auto g = output.ptr<FramePixel>(i);
		auto p = src.ptr<uchar>(i);
		g[0].gx = (p[1] - p[0]) / 2;
		g[cols - 1].gx = (p[cols - 1] - p[cols - 2]) / 2;
	}
}

void computeGradient(const Mat &_src, Mat &gx, Mat &gy)
{
	Mat src;
	if (_src.type() == CV_8UC3)
	{
		cvtColor(_src, src, CV_BGR2GRAY);
	}
	else
	{
		src = _src;
	}
	gx.create(src.size(), CV_32F);
	gy.create(src.size(), CV_32F);
	for (size_t i = 1; i < src.rows - 1; i++)
	{
		auto p0 = src.ptr<uchar>(i - 1);
		auto p1 = src.ptr<uchar>(i);
		auto p2 = src.ptr<uchar>(i + 1);
		auto g = gy.ptr<float>(i);
		for (size_t j = 0; j < src.cols; j++)
		{
			g[j] = (p2[j] - p0[j]) / 2;
		}
	}
	auto g = gy.ptr<float>(0);
	auto g2 = gy.ptr<float>(gy.rows - 1);
	auto p0 = src.ptr<uchar>(0);
	auto p1 = src.ptr<uchar>(1);
	auto p2 = src.ptr<uchar>(src.rows - 2);
	auto p3 = src.ptr<uchar>(src.rows - 1);
	for (size_t j = 0; j < src.cols; j++)
	{
		g[j] = (p1[j] - p0[j]) / 2;
		g2[j] = (p3[j] - p2[j]) / 2;
	}

	for (size_t i = 0; i < src.rows; i++)
	{
		auto p0 = src.ptr<uchar>(i);
		auto g = gx.ptr<float>(i);
		for (size_t j = 1; j < src.cols - 1; j++)
		{
			g[j] = (p0[j + 1] - p0[j - 1]) / 2;
		}
	}

	for (size_t i = 0; i < src.rows; i++)
	{
		const int cols = gx.cols;
		auto g = gx.ptr<float>(i);
		auto p = src.ptr<uchar>(i);
		g[0] = (p[1] - p[0]) / 2;
		g[cols - 1] = (p[cols - 1] - p[cols - 2]) / 2;
	}
}

void computeJacobian(vector<MarkerPixel> &M, const superFrame &I, IntrinsicMatrix&K, Camera &C, const float*p, const vector<Point3f> &kp3, const vector<Point2f> &kp2, double w1, double w2, Mat &J)
{
	Point3f c;
	Point2f p_i;
	const auto &tt = (float*)C.t.data;
	double rx, ry, rz;
	J = Mat_<double>(1, 6, 0.0);
	auto j = (double*)J.data;
	Mat J2 = Mat_<double>(1, 6);
	auto j2 = (double*)J2.data;
	double e0;
	for (size_t i = 0; i < M.size(); i++)
	{
		auto &m = M[i];
		C.TransformToCameraPlane(m.m_3d, c);
		K.projectToImage(c, p_i);
		auto & fp = I.ptr<FramePixel>((int)p_i.y)[(int)p_i.x];
		e0 = p[0] * fp.i + p[1] - m.i;
		rx = c.x - tt[0]; ry = c.y - tt[1]; rz = c.z - tt[2];
		j2[3] = 2 * e0*w1*(p[0] * K.fx* fp.gx) / c.z;
		j2[4] = 2 * e0*w1*(p[0] * K.fy *fp.gy) / c.z;
		j2[5] = 2 * e0*w1*(-(p[0] * c.x*K.fx*fp.gx) / (c.z*c.z) - (p[0] * c.y*K.fy*fp.gy) / (c.z*c.z));
		j2[0] = j2[4] * rz - j2[5] * ry;
		j2[1] = j2[5] * rx - j2[3] * rz;
		j2[2] = j2[3] * ry - j2[4] * rx;
		for (size_t ii = 0; ii < 6; ii++)
		{
			j[ii] += j2[ii];
		}
	}
	vector<Point3f> kp3_c;
	C.TransformToCameraPlane(kp3, kp3_c);
	vector<Point2f> kp2_d;
	K.projectToImage(kp3_c, kp2_d);

	for (size_t i = 0; i < kp2_d.size(); i++)
	{
		auto &t = kp2_d[i];
		auto &c = kp3_c[i];
		j2[3] = 2 * w2*(K.fx*t.x) / c.z;
		j2[4] = 2 * w2*(K.fy*t.y) / c.z;
		j2[5] = 2 * w2*(-(c.x*K.fx*t.x) / (c.z*c.z) - (c.y*K.fy*t.y) / (c.z*c.z));
		j2[0] = j2[4] * rz - j2[5] * ry;
		j2[1] = j2[5] * rx - j2[3] * rz;
		j2[2] = j2[3] * ry - j2[4] * rx;
		for (size_t ii = 0; ii < 6; ii++)
		{
			j[ii] += j2[ii];
		}
	}

}

